//	Aviv Charikar	Yael Tzirulnikov	305343600	307955609

package main;

import java.io.*;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;

import modelTraining.*;

public class Ex3 {
	public static void main(String[] args) throws Exception {
		if (args.length < 5)
			throw new Exception("Not enough arguments");
		checkModels(args[0], args[1], args[2], args[3], args[4]);
	}
	
	private static void checkModels(String devFilename, String testFilename, String inputWord1, String inputWord2, String outputFilename) throws IOException {
		PrintWriter writer = new PrintWriter(outputFilename);
		
		writer.println("#Students\tAviv Charikar\tYael Tzirulnikov\t305343600\t307955609");
		
		writer.println("#Output1\t" + devFilename);
		writer.println("#Output2\t" + testFilename);
		writer.println("#Output3\t" + inputWord1 + " " + inputWord2);
		writer.println("#Output4\t" + outputFilename);
		
		// The vocabulary size is known
		int vocSize = 300000;
		writer.println("#Output5\t" + vocSize);
		
		// Reads the whole dev file with a basic model
		BasicModelTraining basic = new BasicModelTraining(devFilename);
		int eventsNo = basic.getUnigramCount();
		writer.println("#Output6\t" + eventsNo);
		
		// Lidstone using 90% of the events for training
		int trainSize = (int)Math.round(0.9 * eventsNo);
		LidstoneModelTraining lidstone = new LidstoneModelTraining(vocSize, devFilename, trainSize);
		writer.println("#Output7\t" + lidstone.getCountValidation());
		writer.println("#Output8\t" + lidstone.getCountTraining());
		writer.println("#Output9\t" + lidstone.getAllTrainingWords().size());
		writer.println("#Output10\t" + lidstone.getWordCount(inputWord1));
		writer.println("#Output11\t" + lidstone.getBigramCount(inputWord1, inputWord2));
		
		// Perplexity calculations for Lidstone with bigram lambda = 0.0001, 0.001, 0.1
		lidstone.setLambdas(0.1, 0.0001);
		writer.println("#Output12\t" + lidstone.calcPerplexityOnValidationSet());
		lidstone.setLambdas(0.1, 0.001);
		writer.println("#Output13\t" + lidstone.calcPerplexityOnValidationSet());
		lidstone.setLambdas(0.1, 0.1);
		writer.println("#Output14\t" + lidstone.calcPerplexityOnValidationSet());
		
		// Iterates over all possible bigram lambdas (0.0001 - 0.02) and finds the one with the best perplexity
		lidstone.setLambdas(0.1, 0.0001);
		double bestLambda = 0.0001;
		double bestPerp = lidstone.calcPerplexityOnValidationSet();
		for (double lambda = 0.0002; lambda <= 0.02; lambda += 0.0001) {
			lidstone.setLambdas(0.1, lambda);
			double perp = lidstone.calcPerplexityOnValidationSet();
			if (perp < bestPerp) {
				bestLambda = lambda;
				bestPerp = perp;
			}
		}
		
		// Print the best lambda2 and its perplexity
		lidstone.setLambdas(0.1, bestLambda);
		DecimalFormat df = new DecimalFormat("#.####");
		writer.println("#Output15\t" + df.format(bestLambda));
		writer.println("#Output16\t" + bestPerp);
		
		lidstone.setLambdas(0.1, bestLambda);
		
		// Reads the whole test file with a basic model
		BasicModelTraining testBasic = new BasicModelTraining(testFilename);
		
		// Calculate lidstone model perplexity on test set
		double testLidstonePerp = lidstone.calcPerplexityOnTestSet(testBasic);
		writer.println("#Output17\t" + testLidstonePerp);

		writer.println("#Output18");
		lidstone.setLambdas(0.1, 0.001);
		ArrayList<BackOffRecord> records = new ArrayList<BackOffRecord>();
		
		// Get all training words data
		for (String word : lidstone.getAllTrainingWords()) {
			records.add(new BackOffRecord(word, lidstone.getBigramCount(inputWord1, word), 
										  lidstone.getBigramProbability(inputWord1, word)));
		}
		
		// Sort according to backoff probability
		Collections.sort(records, new Comparator<BackOffRecord>() {
            public int compare(BackOffRecord r1,BackOffRecord r2) {
            	return Double.compare(r2.probability, r1.probability);
            }
        });
		
		// Print the sorted records
		int index = 0;
		for (BackOffRecord record : records) {
			index++;
			writer.print(index + "\t");
			writer.print(record.word  + "\t");
			writer.print(record.bigramCount + "\t");
			writer.println(record.probability);
		}
		
		// Print the line about unseen events
		int numUnseenEvents = vocSize - lidstone.getAllTrainingWords().size();
		writer.print(numUnseenEvents + "\t");
		writer.print("UNSEEN_EVENT\t");
		writer.print("0\t");
		writer.println(lidstone.getUnseenWordProbability());
		
		writer.close();
	}
	
	private static class BackOffRecord {
		public String word;
		public int bigramCount;
		public double probability;
		
		public BackOffRecord(String word, int count, double probability) {
			this.word =  word;
			this.bigramCount = count;
			this.probability = probability;
		}
	}
}
